/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.client;

import static org.sss.util.Util.callInParallel;

import java.net.InetSocketAddress;
import java.rmi.NotBoundException;
import java.rmi.RemoteException;
import java.rmi.registry.LocateRegistry;
import java.rmi.registry.Registry;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import org.slf4j.Logger;
import org.sss.common.SssServerRemote;
import org.sss.mapreduce.Cluster;
import org.sss.mapreduce.Configuration;
import org.sss.mapreduce.SssException;
import org.sss.mapreduce.StorageNode;
import org.sss.util.BinaryCallable;
import org.sss.util.UnaryCallable;

/**
 * Manager of communications to remote nodes.
 *
 * This class manages RMI stubs to communicate remote nodes.
 * TODO: implement re-connect.
 * NOTE: This class is used by only this module.
 */
class ClusterManager {
  @SuppressWarnings("unused")
  private static Logger logger = SssClient.getLogger(ClusterManager.class);
  private Cluster cluster;
  private Set<JobWaiter<?>> jobWaiters = new HashSet<JobWaiter<?>>();
  private KeepAliveTimer keepAliveTimer;
  private Configuration conf;
  private int workerCounter = 0;

  /**
   * RMI stub and lock.
   */
  public static class Server {
    private SssServerRemote rmiObject;
    private Lock token = new ReentrantLock();

    public Server(InetSocketAddress worker) throws SssException {
      try {
        Registry registry =
          LocateRegistry.getRegistry(worker.getHostName(), worker.getPort());
        this.rmiObject = (SssServerRemote) registry.lookup("SssServerImpl");
      } catch (RemoteException e) {
        throw new SssException(e);
      } catch (NotBoundException e) {
        throw new SssException(e);
      }
    }

    public SssServerRemote get() {
      return rmiObject;
    }

    public void lock() {
      token.lock();
    }

    public boolean tryLock() {
      return token.tryLock();
    }

    public void unlock() {
      token.unlock();
    }
  }

  private Map<InetSocketAddress, Server> workerMap = new HashMap<InetSocketAddress, Server>();

  /**
   * create new ClusterManager.
   *
   * @param cluster   Informations of remote nodes.
   * @param conf      Configuration. This class uses the followings.
   *                  <dl> 
   *                    <dt> heartbeat.period </dt> 
   *                    <dd> The interval to which this class sends a 
   *                         heartbeat. The unit is second. Default value 
   *                         is 60.
   *                    </dd> 
   *                  </dl>
   */
  public ClusterManager(Cluster cluster, Configuration conf) {
    this.cluster = cluster;
    this.keepAliveTimer = new KeepAliveTimer(this);
    this.conf = conf;
  }

  /**
   * Get RMI stub and lock.
   *
   * @param node data node.
   * @return The locked RMI stub to related to <code>node</code>.
   */
  public Server get(StorageNode node) throws SssException {
    Server s = getOrConnect(cluster.getWorkerNode(node));
    s.lock();
    return s;
  }

  /**
   * Get any RMI stub and lock.
   *
   * @return any locked RMI stub.
   */
  public synchronized Server getAny() throws SssException {
    int i = Math.abs(workerCounter++) % cluster.getWorkerNodes().size();
    StorageNode sn = cluster.getStorageNodes().get(i);
    Server s = getOrConnect(cluster.getWorkerNode(sn));
    s.lock();
    return s;
  }

  private Server getOrConnect(InetSocketAddress worker) throws SssException {
    if (workerMap.containsKey(worker)) {
      return workerMap.get(worker);
    } else {
      Server server = new Server(worker);
      workerMap.put(worker, server);
      return server;
    }
  }

  /**
   * Get all RMI stubs and address.
   */
  public Map<InetSocketAddress, Server> getServerMap() throws SssException {
    for (InetSocketAddress worker: cluster.getWorkerNodes()) {
      getOrConnect(worker);
    }
    return Collections.unmodifiableMap(workerMap);
  }

  /**
   * Get all RMI stubs.
   */
  public Iterable<Server> getAllServers() throws SssException {
    return getServerMap().values();
  }

  /**
   * Get all storage nodes.
   */
  public List<StorageNode> getStorageNodes() {
    return cluster.getStorageNodes();
  }

  public synchronized void register(JobWaiter<?> waiter) throws SssException {
    boolean start = false;
    if (jobWaiters.isEmpty()) {
      start = true;
    }
    jobWaiters.add(waiter);

    if (start) {
      keepAliveTimer.start(conf.getInt("heartbeat.period", 60));
    }
  }

  public synchronized void unregister(JobWaiter<?> waiter) {
    jobWaiters.remove(waiter);
    if (jobWaiters.isEmpty()) {
      keepAliveTimer.stop();
    }
  }

  /**
   * Receive keep alive error.
   *
   * NOTE: This used by only @{link org.sss.client.KeepAliveTimer}.
   */
  synchronized void notifyKeepAliveError(Throwable e) {
    for (JobWaiter<?> waiter: jobWaiters) {
      waiter.setKeepAliveError(e);
    }
  }
  
  /**
   * .
   * TODO: does not thread happy?
   */
  public <T> List<T> callEachServerInParallel(final UnaryCallable<T, SssServerRemote> callback) throws SssException {
    try {
      return callInParallel(new UnaryCallable<T, Server>() {
        @Override
        public T call(Server server) throws Exception {
          server.lock();
          try {
            return callback.call(server.get());
          }
          finally {
            server.unlock();
          }
        }
      }, getAllServers());
    } catch (InterruptedException e) {
      throw new SssException(e);
    } catch (ExecutionException e) {
      Throwable ee = e.getCause();
      if (ee != null) {
        throw new SssException(ee);
      }
      else {
        throw new SssException(e);
      }
    }
  }
  
  public <T> List<T> callEachStorageInParallel(
      final BinaryCallable<T, SssServerRemote, StorageNode> callback
      ) throws SssException {
    try {
      return callInParallel(new UnaryCallable<T, StorageNode>() {
        @Override
        public T call(StorageNode sn) throws Exception {
          Server server = get(sn);
          try {
            return callback.call(server.get(), sn);
          }
          finally {
            server.unlock();
          }
        }
      }, getStorageNodes());
    } catch (InterruptedException e) {
      throw new SssException(e);
    } catch (ExecutionException e) {
      Throwable ee = e.getCause();
      if (ee != null) {
        throw new SssException(ee);
      }
      else {
        throw new SssException(e);
      }
    }
  }
}
